package data

import (
	"fmt"
	"net"
	"rtps/common"
	"rtps/discovery"
	"rtps/messages"
	"rtps/transport"
	"sync"
	"time"
)

// 数据Writer
type DataWriter struct {
	WriterEndpoint *discovery.WriterEndpoint
	Transport      *transport.UDPTransport
	SequenceNumber uint64
	Readers        map[string]*ReaderInfo // key: Reader GUID string
	mutex          sync.RWMutex
}

// Reader信息
type ReaderInfo struct {
	ReaderGUID   common.GUID
	Locator      common.Locator
	LastAckNack  time.Time
	LastSequence uint64
	Reliable     bool
}

// 数据Reader
type DataReader struct {
	ReaderEndpoint *discovery.ReaderEndpoint
	Transport      *transport.UDPTransport
	Writers        map[string]*WriterInfo // key: Writer GUID string
	ReceivedData   map[uint64][]byte      // key: Sequence number
	LastSequence   uint64
	mutex          sync.RWMutex
}

// Writer信息
type WriterInfo struct {
	WriterGUID    common.GUID
	Locator       common.Locator
	LastHeartbeat time.Time
	LastSequence  uint64
	Reliable      bool
}

// 创建新的DataWriter
func NewDataWriter(writerEndpoint *discovery.WriterEndpoint, trans *transport.UDPTransport) *DataWriter {
	return &DataWriter{
		WriterEndpoint: writerEndpoint,
		Transport:      trans,
		SequenceNumber: 1,
		Readers:        make(map[string]*ReaderInfo),
	}
}

// 创建新的DataReader
func NewDataReader(readerEndpoint *discovery.ReaderEndpoint, trans *transport.UDPTransport) *DataReader {
	return &DataReader{
		ReaderEndpoint: readerEndpoint,
		Transport:      trans,
		Writers:        make(map[string]*WriterInfo),
		ReceivedData:   make(map[uint64][]byte),
		LastSequence:   0,
	}
}

// 写入数据
func (dw *DataWriter) Write(data []byte) error {
	dw.mutex.Lock()
	defer dw.mutex.Unlock()

	// 创建数据消息
	msg := messages.NewMessage(dw.WriterEndpoint.GUID.Prefix)

	// 添加Data子消息
	dataSubmsg := &messages.Data{
		Header: messages.SubmessageHeader{
			SubmessageID: common.MessageTypeData,
			Flags:        0x01, // 关键数据标志
		},
		ReaderID: messages.EntityID{
			EntityKey:  [3]byte{0x00, 0x00, 0x00}, // 通配Reader
			EntityKind: 0xC7,                      // Reader实体类型
		},
		WriterID: messages.EntityID{
			EntityKey:  [3]byte{dw.WriterEndpoint.GUID.EntityID[0], dw.WriterEndpoint.GUID.EntityID[1], dw.WriterEndpoint.GUID.EntityID[2]},
			EntityKind: 0xC2, // Writer实体类型
		},
		WriterSN: messages.SequenceNumber{
			High: 0,
			Low:  uint32(dw.SequenceNumber),
		},
		SerializedData: data,
	}
	msg.AddSubmessage(dataSubmsg)

	// 发送给所有匹配的Reader
	for _, reader := range dw.Readers {
		addr := &net.UDPAddr{
			IP:   bytesToIP(reader.Locator.Address),
			Port: int(reader.Locator.Port),
		}

		if err := dw.Transport.SendMessage(msg, addr); err != nil {
			fmt.Printf("Failed to send data to reader %s: %v\n",
				guidToString(reader.ReaderGUID), err)
		}
	}

	dw.SequenceNumber++
	return nil
}

// 读取数据
func (dr *DataReader) Read() ([]byte, error) {
	dr.mutex.Lock()
	defer dr.mutex.Unlock()

	// 检查是否有可用的数据
	if len(dr.ReceivedData) == 0 {
		return nil, fmt.Errorf("no data available")
	}

	// 按顺序返回数据
	for seq := dr.LastSequence + 1; ; seq++ {
		if data, exists := dr.ReceivedData[seq]; exists {
			delete(dr.ReceivedData, seq)
			dr.LastSequence = seq
			return data, nil
		}
		break
	}

	return nil, fmt.Errorf("no new data available")
}

// 处理接收到的数据消息
func (dr *DataReader) HandleDataMessage(msg *messages.Message, addr *net.UDPAddr) {
	for _, submsg := range msg.Submessages {
		switch sm := submsg.(type) {
		case *messages.Data:
			// 检查是否是匹配的Writer
			writerGUID := common.GUID{
				Prefix:   msg.Header.GUIDPrefix,
				EntityID: [4]byte{sm.WriterID.EntityKey[0], sm.WriterID.EntityKey[1], sm.WriterID.EntityKey[2], sm.WriterID.EntityKind},
			}

			if dr.isMatchingWriter(writerGUID) {
				dr.processData(sm, writerGUID, addr)
			}
		case *messages.Heartbeat:
			// 处理心跳消息
			writerGUID := common.GUID{
				Prefix:   msg.Header.GUIDPrefix,
				EntityID: [4]byte{sm.WriterID.EntityKey[0], sm.WriterID.EntityKey[1], sm.WriterID.EntityKey[2], sm.WriterID.EntityKind},
			}

			if dr.isMatchingWriter(writerGUID) {
				dr.processHeartbeat(sm, writerGUID, addr)
			}
		}
	}
}

// 处理数据消息
func (dr *DataReader) processData(data *messages.Data, writerGUID common.GUID, addr *net.UDPAddr) {
	dr.mutex.Lock()
	defer dr.mutex.Unlock()

	sequence := uint64(data.WriterSN.Low)

	// 存储接收到的数据
	dr.ReceivedData[sequence] = data.SerializedData

	// 更新Writer信息
	writerInfo, exists := dr.Writers[guidToString(writerGUID)]
	if exists {
		writerInfo.LastSequence = sequence
	}

	// 发送AckNack（如果是可靠传输）
	if dr.ReaderEndpoint.QoS.Reliability == common.QoSPolicyReliabilityReliable {
		dr.sendAckNack(writerGUID, addr, sequence)
	}
}

// 处理心跳消息
func (dr *DataReader) processHeartbeat(heartbeat *messages.Heartbeat, writerGUID common.GUID, addr *net.UDPAddr) {
	dr.mutex.Lock()
	defer dr.mutex.Unlock()

	writerInfo, exists := dr.Writers[guidToString(writerGUID)]
	if exists {
		writerInfo.LastHeartbeat = time.Now()
		writerInfo.LastSequence = uint64(heartbeat.LastSN.Low)

		// 检查是否有丢失的数据
		if dr.ReaderEndpoint.QoS.Reliability == common.QoSPolicyReliabilityReliable {
			dr.checkForMissingData(writerGUID, addr, heartbeat)
		}
	}
}

// 检查丢失的数据并发送AckNack
func (dr *DataReader) checkForMissingData(writerGUID common.GUID, addr *net.UDPAddr, heartbeat *messages.Heartbeat) {
	firstSeq := uint64(heartbeat.FirstSN.Low)
	lastSeq := uint64(heartbeat.LastSN.Low)

	var missingSequences []uint64

	for seq := firstSeq; seq <= lastSeq; seq++ {
		if _, exists := dr.ReceivedData[seq]; !exists {
			missingSequences = append(missingSequences, seq)
		}
	}

	if len(missingSequences) > 0 {
		dr.sendAckNackWithMissing(writerGUID, addr, firstSeq, lastSeq, missingSequences)
	}
}

// 发送AckNack消息
func (dr *DataReader) sendAckNack(writerGUID common.GUID, addr *net.UDPAddr, lastSequence uint64) {
	msg := messages.NewMessage(dr.ReaderEndpoint.GUID.Prefix)

	ackNack := &messages.AckNack{
		Header: messages.SubmessageHeader{
			SubmessageID: common.MessageTypeAckNack,
			Flags:        0x01,
		},
		ReaderID: messages.EntityID{
			EntityKey:  [3]byte{dr.ReaderEndpoint.GUID.EntityID[0], dr.ReaderEndpoint.GUID.EntityID[1], dr.ReaderEndpoint.GUID.EntityID[2]},
			EntityKind: 0xC7,
		},
		WriterID: messages.EntityID{
			EntityKey:  [3]byte{writerGUID.EntityID[0], writerGUID.EntityID[1], writerGUID.EntityID[2]},
			EntityKind: 0xC2,
		},
		ReaderSN: messages.SequenceNumber{
			High: 0,
			Low:  uint32(lastSequence),
		},
		Count: 1,
	}

	msg.AddSubmessage(ackNack)

	if err := dr.Transport.SendMessage(msg, addr); err != nil {
		fmt.Printf("Failed to send AckNack to writer %s: %v\n",
			guidToString(writerGUID), err)
	}
}

// 发送包含丢失序列的AckNack消息
func (dr *DataReader) sendAckNackWithMissing(writerGUID common.GUID, addr *net.UDPAddr,
	firstSeq, lastSeq uint64, missingSequences []uint64) {
	msg := messages.NewMessage(dr.ReaderEndpoint.GUID.Prefix)

	// 创建位图表示丢失的序列
	bitmap := createBitmap(firstSeq, lastSeq, missingSequences)

	ackNack := &messages.AckNack{
		Header: messages.SubmessageHeader{
			SubmessageID: common.MessageTypeAckNack,
			Flags:        0x01,
		},
		ReaderID: messages.EntityID{
			EntityKey:  [3]byte{dr.ReaderEndpoint.GUID.EntityID[0], dr.ReaderEndpoint.GUID.EntityID[1], dr.ReaderEndpoint.GUID.EntityID[2]},
			EntityKind: 0xC7,
		},
		WriterID: messages.EntityID{
			EntityKey:  [3]byte{writerGUID.EntityID[0], writerGUID.EntityID[1], writerGUID.EntityID[2]},
			EntityKind: 0xC2,
		},
		ReaderSN: messages.SequenceNumber{
			High: 0,
			Low:  uint32(lastSeq),
		},
		Count:  uint32(len(missingSequences)),
		Bitmap: bitmap,
	}

	msg.AddSubmessage(ackNack)

	if err := dr.Transport.SendMessage(msg, addr); err != nil {
		fmt.Printf("Failed to send AckNack with bitmap to writer %s: %v\n",
			guidToString(writerGUID), err)
	}
}

// 定期发送心跳消息
func (dw *DataWriter) StartHeartbeat(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				dw.sendHeartbeat()
			}
		}
	}()
}

// 发送心跳消息
func (dw *DataWriter) sendHeartbeat() {
	dw.mutex.RLock()
	defer dw.mutex.RUnlock()

	if len(dw.Readers) == 0 {
		return
	}

	msg := messages.NewMessage(dw.WriterEndpoint.GUID.Prefix)

	heartbeat := &messages.Heartbeat{
		Header: messages.SubmessageHeader{
			SubmessageID: common.MessageTypeHeartbeat,
			Flags:        0x01,
		},
		ReaderID: messages.EntityID{
			EntityKey:  [3]byte{0x00, 0x00, 0x00}, // 通配Reader
			EntityKind: 0xC7,
		},
		WriterID: messages.EntityID{
			EntityKey:  [3]byte{dw.WriterEndpoint.GUID.EntityID[0], dw.WriterEndpoint.GUID.EntityID[1], dw.WriterEndpoint.GUID.EntityID[2]},
			EntityKind: 0xC2,
		},
		FirstSN: messages.SequenceNumber{
			High: 0,
			Low:  1,
		},
		LastSN: messages.SequenceNumber{
			High: 0,
			Low:  uint32(dw.SequenceNumber - 1),
		},
		Count: 1,
	}

	msg.AddSubmessage(heartbeat)

	// 发送给所有Reader
	for _, reader := range dw.Readers {
		addr := &net.UDPAddr{
			IP:   bytesToIP(reader.Locator.Address),
			Port: int(reader.Locator.Port),
		}

		if err := dw.Transport.SendMessage(msg, addr); err != nil {
			fmt.Printf("Failed to send heartbeat to reader %s: %v\n",
				guidToString(reader.ReaderGUID), err)
		}
	}
}

// 处理AckNack消息
func (dw *DataWriter) HandleAckNack(msg *messages.Message, addr *net.UDPAddr) {
	for _, submsg := range msg.Submessages {
		switch sm := submsg.(type) {
		case *messages.AckNack:
			readerGUID := common.GUID{
				Prefix:   msg.Header.GUIDPrefix,
				EntityID: [4]byte{sm.ReaderID.EntityKey[0], sm.ReaderID.EntityKey[1], sm.ReaderID.EntityKey[2], sm.ReaderID.EntityKind},
			}

			if dw.isMatchingReader(readerGUID) {
				dw.processAckNack(sm, readerGUID, addr)
			}
		}
	}
}

// 处理AckNack消息
func (dw *DataWriter) processAckNack(ackNack *messages.AckNack, readerGUID common.GUID, addr *net.UDPAddr) {
	dw.mutex.Lock()
	defer dw.mutex.Unlock()

	readerInfo, exists := dw.Readers[guidToString(readerGUID)]
	if exists {
		readerInfo.LastAckNack = time.Now()

		// 处理丢失的数据重传
		if ackNack.Count > 0 && len(ackNack.Bitmap) > 0 {
			dw.retransmitMissingData(readerGUID, addr, ackNack)
		}
	}
}

// 重传丢失的数据
func (dw *DataWriter) retransmitMissingData(readerGUID common.GUID, addr *net.UDPAddr, ackNack *messages.AckNack) {
	// 这里需要实现数据缓存和重传逻辑
	// 在实际实现中，Writer需要缓存已发送的数据以便重传
	fmt.Printf("Need to retransmit data for reader %s\n", guidToString(readerGUID))
}

// 检查是否匹配的Reader
func (dw *DataWriter) isMatchingReader(readerGUID common.GUID) bool {
	// 检查Reader是否订阅了Writer的主题
	// 这里需要实现主题匹配逻辑
	return true // 简化实现
}

// 检查是否匹配的Writer
func (dr *DataReader) isMatchingWriter(writerGUID common.GUID) bool {
	// 检查Writer是否发布了Reader感兴趣的主题
	// 这里需要实现主题匹配逻辑
	return true // 简化实现
}

// 添加匹配的Reader
func (dw *DataWriter) AddReader(readerGUID common.GUID, locator common.Locator, reliable bool) {
	dw.mutex.Lock()
	defer dw.mutex.Unlock()

	dw.Readers[guidToString(readerGUID)] = &ReaderInfo{
		ReaderGUID:  readerGUID,
		Locator:     locator,
		LastAckNack: time.Now(),
		Reliable:    reliable,
	}
}

// 添加匹配的Writer
func (dr *DataReader) AddWriter(writerGUID common.GUID, locator common.Locator, reliable bool) {
	dr.mutex.Lock()
	defer dr.mutex.Unlock()

	dr.Writers[guidToString(writerGUID)] = &WriterInfo{
		WriterGUID:    writerGUID,
		Locator:       locator,
		LastHeartbeat: time.Now(),
		Reliable:      reliable,
	}
}

// 辅助函数
func guidToString(guid common.GUID) string {
	return fmt.Sprintf("%x%x", guid.Prefix, guid.EntityID)
}

func bytesToIP(bytes [16]byte) net.IP {
	if bytes[0] == 0 && bytes[1] == 0 && bytes[2] == 0 && bytes[3] == 0 &&
		bytes[4] == 0 && bytes[5] == 0 && bytes[6] == 0 && bytes[7] == 0 &&
		bytes[8] == 0 && bytes[9] == 0 && bytes[10] == 0 && bytes[11] == 0 {
		return net.IPv4(bytes[12], bytes[13], bytes[14], bytes[15])
	}
	return net.IP(bytes[:])
}

func createBitmap(firstSeq, lastSeq uint64, missingSequences []uint64) []byte {
	bitmapSize := (lastSeq - firstSeq + 7) / 8
	bitmap := make([]byte, bitmapSize)

	for _, seq := range missingSequences {
		if seq >= firstSeq && seq <= lastSeq {
			bitIndex := seq - firstSeq
			byteIndex := bitIndex / 8
			bitPos := bitIndex % 8
			bitmap[byteIndex] |= 1 << bitPos
		}
	}

	return bitmap
}
